//! 简单格式化器实现

use crate::model::{AstNode, TokenKind};

/// 格式化器状态
pub struct Formatter {
    /// 输出缓冲区
    pub output: String,
    /// 当前缩进级别
    pub indent_level: usize,
    /// 是否需要在下一行添加缩进
    pub needs_indent: bool,
}

impl Formatter {
    /// 创建新的格式化器
    pub fn new() -> Self {
        Self {
            output: String::new(),
            indent_level: 0,
            needs_indent: false,
        }
    }

    /// 添加缩进
    fn add_indent(&mut self) {
        for _ in 0..self.indent_level {
            self.output.push_str("    "); // 4个空格缩进
        }
        self.needs_indent = false;
    }

    /// 添加字符串到输出
    fn add_str(&mut self, s: &str) {
        if self.needs_indent {
            self.add_indent();
        }
        self.output.push_str(s);
    }

    /// 添加空格
    fn add_space(&mut self) {
        self.output.push(' ');
    }

    /// 添加换行
    fn add_newline(&mut self) {
        self.output.push('\n');
        self.needs_indent = true;
    }

    /// 格式化编译单元
    pub fn format_comp_unit(&mut self, node: &AstNode) -> Result<(), String> {
        if let AstNode::Sequence(items) = node {
            let mut first_func = true;
            
            for (i, item) in items.iter().enumerate() {
                // 函数定义前加空行（除非是第一个函数）
                if let AstNode::Node { kind, .. } = item {
                    if *kind == TokenKind::Int || *kind == TokenKind::Void {
                        if !first_func {
                            self.add_newline();
                        }
                        first_func = false;
                    }
                }
                
                self.format_node(item)?;
                
                // 声明语句之间不加空行
                if i < items.len() - 1 {
                    if let AstNode::Node { kind, .. } = item {
                        if *kind != TokenKind::Int && *kind != TokenKind::Const {
                            self.add_newline();
                        }
                    }
                }
            }
            Ok(())
        } else {
            Err("CompUnit should be a Sequence".to_string())
        }
    }

    /// 格式化节点
    pub fn format_node(&mut self, node: &AstNode) -> Result<(), String> {
        match node {
            AstNode::Token(kind, text_opt) => self.format_token(*kind, text_opt.as_ref()),
            AstNode::Sequence(nodes) => self.format_sequence(nodes),
            AstNode::Node { kind, children, .. } => self.format_complex_node(*kind, children),
        }
    }

    /// 格式化Token
    fn format_token(&mut self, kind: TokenKind, text_opt: Option<&String>) -> Result<(), String> {
        let text = match kind {
            TokenKind::Const => "const",
            TokenKind::Int => "int",
            TokenKind::Void => "void",
            TokenKind::If => "if",
            TokenKind::Else => "else",
            TokenKind::While => "while",
            TokenKind::Break => "break",
            TokenKind::Continue => "continue",
            TokenKind::Return => "return",
            TokenKind::Ident => {
                // 使用原始文本或占位符
                if let Some(text) = text_opt {
                    text.as_str()
                } else {
                    "ident"
                }
            }
            TokenKind::IntegerConst => {
                // 数字常量直接使用原始文本
                if let Some(text) = text_opt {
                    text.as_str()
                } else {
                    "123"
                }
            }
            TokenKind::Plus => "+",
            TokenKind::Minus => "-",
            TokenKind::Mul => "*",
            TokenKind::Div => "/",
            TokenKind::Mod => "%",
            TokenKind::Assign => "=",
            TokenKind::Eq => "==",
            TokenKind::Neq => "!=",
            TokenKind::Lt => "<",
            TokenKind::Le => "<=",
            TokenKind::Gt => ">",
            TokenKind::Ge => ">=",
            TokenKind::Not => "!",
            TokenKind::And => "&&",
            TokenKind::Or => "||",
            TokenKind::LParen => "(",
            TokenKind::RParen => ")",
            TokenKind::LBrack => "[",
            TokenKind::RBrack => "]",
            TokenKind::LBrace => "{",
            TokenKind::RBrace => "}",
            TokenKind::Comma => ",",
            TokenKind::Semicolon => ";",
        };

        self.add_str(text);
        Ok(())
    }

    /// 格式化序列
    fn format_sequence(&mut self, nodes: &[AstNode]) -> Result<(), String> {
        for (i, node) in nodes.iter().enumerate() {
            self.format_node(node)?;
            
            // 在分号后添加换行（语句结束）
            if let AstNode::Token(TokenKind::Semicolon, _) = node {
                self.add_newline();
            }
            
            // 在逗号后加空格（用于数组初始化列表等）
            if let AstNode::Token(TokenKind::Comma, _) = node {
                self.add_space();
            }
            
            // 在右花括号后添加换行
            if let AstNode::Token(TokenKind::RBrace, _) = node {
                self.add_newline();
            }
            
            // 确保语句之间正确换行（处理 print(a);if (cond) 的情况）
            if i < nodes.len() - 1 {
                // 检查下一个节点是否是语句开始
                let next_is_stmt_start = match &nodes[i + 1] {
                    AstNode::Node { kind, .. } => {
                        matches!(*kind, TokenKind::If | TokenKind::While | TokenKind::Return | 
                                TokenKind::Break | TokenKind::Continue)
                    }
                    AstNode::Token(TokenKind::Ident, _) => {
                        // 如果当前是分号，下一个是标识符（可能是函数调用），也添加换行
                        if let AstNode::Token(TokenKind::Semicolon, _) = node {
                            true
                        } else {
                            false
                        }
                    }
                    _ => false,
                };
                
                if next_is_stmt_start {
                    self.add_newline();
                }
            }
        }
        Ok(())
    }

    /// 格式化函数参数序列（确保在同一行）
    fn format_func_params(&mut self, nodes: &[AstNode]) -> Result<(), String> {
        for (i, node) in nodes.iter().enumerate() {
            // 特殊处理Int节点（函数参数）
            if let AstNode::Node { kind, children, .. } = node {
                if *kind == TokenKind::Int {
                    // 格式化参数类型和名称，但不添加换行
                    self.format_token(TokenKind::Int, None)?;
                    self.add_space();
                    if let Some(AstNode::Token(TokenKind::Ident, text_opt)) = children.get(1) {
                        self.format_token(TokenKind::Ident, text_opt.as_ref())?;
                    }
                    continue;
                }
            }
            
            // 处理逗号
            if let AstNode::Token(TokenKind::Comma, _) = node {
                self.format_token(TokenKind::Comma, None)?;
                if i < nodes.len() - 1 {
                    self.add_space();
                }
                continue;
            }
            
            // 默认处理其他节点
            self.format_node(node)?;
        }
        Ok(())
    }

    /// 格式化复杂节点
    fn format_complex_node(&mut self, kind: TokenKind, children: &[AstNode]) -> Result<(), String> {
        match kind {
            TokenKind::Const => self.format_const_decl(children),
            TokenKind::Int => {
                // 检查是否在函数参数上下文中（避免在参数声明后添加换行）
                if self.is_in_func_param_context() {
                    self.format_func_param(children)
                } else {
                    self.format_var_decl(children)
                }
            },
            TokenKind::Void if children.len() >= 5 => self.format_func_def(children),
            TokenKind::LBrace => {
                // 区分数组初始化列表和代码块
                if self.is_init_val_list(children) {
                    self.format_init_val_list(children)
                } else {
                    self.format_block(children)
                }
            }
            TokenKind::If => self.format_if_stmt(children),
            TokenKind::While => self.format_while_stmt(children),
            TokenKind::Return => self.format_return_stmt(children),
            TokenKind::Break => self.format_break_stmt(children),
            TokenKind::Continue => self.format_continue_stmt(children),
            TokenKind::Ident => self.format_lval_or_call(children),
            _ => self.format_expression(children),
        }
    }

    /// 格式化常量声明
    fn format_const_decl(&mut self, children: &[AstNode]) -> Result<(), String> {
        // ConstDecl → 'const' 'int' ConstDef {',' ConstDef} ';'
        for (_i, child) in children.iter().enumerate() {
            self.format_node(child)?;
            
            // 在关键字后加空格
            if let AstNode::Token(kind, _) = child {
                if *kind == TokenKind::Const || *kind == TokenKind::Int {
                    self.add_space();
                }
            }
            
            // 逗号后加空格
            if let AstNode::Token(TokenKind::Comma, _) = child {
                self.add_space();
            }
        }
        self.add_newline();
        Ok(())
    }

    /// 格式化变量声明
    fn format_var_decl(&mut self, children: &[AstNode]) -> Result<(), String> {
        // VarDecl → 'int' VarDef {',' VarDef} ';'
        for (i, child) in children.iter().enumerate() {
            self.format_node(child)?;
            
            // 在int后加空格
            if i == 0 {
                self.add_space();
            }
            
            // 在赋值运算符前加空格
            if let AstNode::Token(TokenKind::Assign, _) = child {
                self.add_space();
            }
            
            // 在赋值运算符后加空格
            if i < children.len() - 1 {
                if let AstNode::Token(TokenKind::Assign, _) = child {
                    if let Some(next_child) = children.get(i + 1) {
                        // 确保赋值运算符后面有空格
                        if !matches!(next_child, 
                            AstNode::Token(TokenKind::Semicolon | TokenKind::Comma, _)) {
                            self.add_space();
                        }
                    }
                }
            }
            
            // 逗号后加空格
            if let AstNode::Token(TokenKind::Comma, _) = child {
                self.add_space();
            }
        }
        self.add_newline();
        Ok(())
    }

    /// 格式化函数定义
    fn format_func_def(&mut self, children: &[AstNode]) -> Result<(), String> {
        // FuncDef → FuncType Ident '(' [FuncFParams] ')' Block
        // 手动处理各个部分以确保正确的格式
        
        // 1. 返回类型
        self.format_node(&children[0])?; // FuncType (void/int)
        self.add_space();
        
        // 2. 函数名
        self.format_node(&children[1])?; // Ident
        
        // 3. 左括号
        self.format_node(&children[2])?; // LParen
        
        // 4. 函数参数序列（如果有）
        if let AstNode::Sequence(nodes) = &children[3] {
            self.format_func_params(nodes)?;
        } else {
            self.format_node(&children[3])?;
        }
        
        // 5. 右括号（确保和参数在同一行）
        self.format_node(&children[4])?; // RParen
        
        // 6. 在右括号和左花括号之间加空格
        self.add_space();
        
        // 7. 代码块
        self.format_node(&children[5])?; // Block
        
        self.add_newline();
        Ok(())
    }

    /// 格式化代码块
    fn format_block(&mut self, children: &[AstNode]) -> Result<(), String> {
        // Block → '{' { BlockItem } '}'
        self.add_str("{");
        self.add_newline();
        
        self.indent_level += 1;
        
        // 格式化块内内容
        for child in children.iter().skip(1).take(children.len() - 2) {
            self.format_node(child)?;
        }
        
        self.indent_level -= 1;
        self.add_str("}");
        self.add_newline();
        
        Ok(())
    }

    /// 格式化if语句
    fn format_if_stmt(&mut self, children: &[AstNode]) -> Result<(), String> {
        // if (cond) stmt [else stmt]
        self.format_node(&children[0])?; // if
        self.add_space();
        self.format_node(&children[1])?; // (
        self.format_node(&children[2])?; // cond
        self.format_node(&children[3])?; // )
        self.add_space();
        self.format_node(&children[4])?; // stmt
        
        if children.len() > 5 {
            // else 语句在新行上，并确保正确缩进
            self.add_newline();
            
            // 保存当前缩进状态
            let current_indent = self.indent_level;
            
            // 格式化else关键字
            self.format_node(&children[5])?; // else
            
            // 在else和左花括号之间加空格
            if children.len() > 6 {
                if let AstNode::Node { kind, .. } = &children[6] {
                    if *kind == TokenKind::LBrace {
                        self.add_space();
                    }
                }
                
                // 如果else后面是if（else if），则加空格
                if let AstNode::Token(TokenKind::If, _) = &children[6] {
                    self.add_space();
                }
                self.format_node(&children[6])?; // else stmt 或 if
            } else if children.len() == 6 {
                // 处理 else { 的情况，确保有空格
                if let AstNode::Node { kind, .. } = &children[5] {
                    if *kind == TokenKind::LBrace {
                        self.add_space();
                    }
                }
            }
            
            // 恢复缩进级别
            self.indent_level = current_indent;
        }
        
        Ok(())
    }

    /// 格式化while语句
    fn format_while_stmt(&mut self, children: &[AstNode]) -> Result<(), String> {
        // while (cond) stmt
        self.format_node(&children[0])?; // while
        self.add_space();
        self.format_node(&children[1])?; // (
        self.format_node(&children[2])?; // cond
        self.format_node(&children[3])?; // )
        self.add_space();
        self.format_node(&children[4])?; // stmt
        
        Ok(())
    }

    /// 格式化return语句
    fn format_return_stmt(&mut self, children: &[AstNode]) -> Result<(), String> {
        // return [exp] ;
        self.format_node(&children[0])?; // return
        
        if children.len() > 1 && !matches!(children[1], AstNode::Token(TokenKind::Semicolon, _)) {
            self.add_space();
            self.format_node(&children[1])?; // exp
        }
        
        // return后直接跟分号不需要空格
        if let Some(AstNode::Token(TokenKind::Semicolon, _)) = children.last() {
            self.format_node(children.last().unwrap())?;
        }
        
        self.add_newline();
        Ok(())
    }

    /// 格式化break语句
    fn format_break_stmt(&mut self, children: &[AstNode]) -> Result<(), String> {
        // break ;
        self.format_node(&children[0])?; // break
        self.format_node(&children[1])?; // ;
        self.add_newline();
        Ok(())
    }

    /// 格式化continue语句
    fn format_continue_stmt(&mut self, children: &[AstNode]) -> Result<(), String> {
        // continue ;
        self.format_node(&children[0])?; // continue
        self.format_node(&children[1])?; // ;
        self.add_newline();
        Ok(())
    }

    /// 格式化左值或函数调用
    fn format_lval_or_call(&mut self, children: &[AstNode]) -> Result<(), String> {
        // 简单实现：直接格式化所有子节点
        for child in children {
            self.format_node(child)?;
        }
        Ok(())
    }

    /// 格式化表达式
    fn format_expression(&mut self, children: &[AstNode]) -> Result<(), String> {
        for (i, child) in children.iter().enumerate() {
            // 处理二元运算符前的空格
            if i > 0 {
                if let AstNode::Token(kind, _) = child {
                    if self.is_binary_operator(*kind) {
                        self.add_space();
                    }
                }
            }
            
            self.format_node(child)?;
            
            // 处理二元运算符后的空格
            if i < children.len() - 1 {
                if let AstNode::Token(kind, _) = child {
                    if self.is_binary_operator(*kind) {
                        self.add_space();
                    }
                }
            }
            
            // 在逗号后加空格
            if let AstNode::Token(TokenKind::Comma, _) = child {
                self.add_space();
            }
            
            // 处理赋值运算符的特殊情况（确保 a=1 变成 a = 1）
            if let AstNode::Token(TokenKind::Assign, _) = child {
                // 如果赋值运算符后面不是空格，添加空格
                if i < children.len() - 1 {
                    if let Some(next_child) = children.get(i + 1) {
                        // 检查下一个token不是运算符或括号
                        if !matches!(next_child, 
                            AstNode::Token(TokenKind::Plus | TokenKind::Minus | 
                                         TokenKind::Mul | TokenKind::Div | 
                                         TokenKind::LParen | TokenKind::RParen, _)) {
                            self.add_space();
                        }
                    }
                }
            }
        }
        Ok(())
    }

    /// 判断是否为初始化值列表（数组初始化）
    fn is_init_val_list(&self, children: &[AstNode]) -> bool {
        // 初始化值列表包含逗号Token
        children.iter().any(|child| {
            if let AstNode::Token(TokenKind::Comma, _) = child {
                true
            } else {
                false
            }
        })
    }

    /// 格式化初始化值列表（数组初始化）
    fn format_init_val_list(&mut self, children: &[AstNode]) -> Result<(), String> {
        // 初始化值列表格式：{ init_val, init_val, ... }
        self.add_str("{");
        
        // 格式化内部内容（跳过第一个{和最后一个}）
        for (_i, child) in children.iter().skip(1).take(children.len() - 2).enumerate() {
            self.format_node(child)?;
            
            // 在逗号后加空格
            if let AstNode::Token(TokenKind::Comma, _) = child {
                self.add_space();
            }
        }
        
        self.add_str("}");
        Ok(())
    }

    /// 判断是否为二元运算符
    fn is_binary_operator(&self, kind: TokenKind) -> bool {
        matches!(
            kind,
            TokenKind::Plus
                | TokenKind::Minus
                | TokenKind::Mul
                | TokenKind::Div
                | TokenKind::Mod
                | TokenKind::Assign
                | TokenKind::Eq
                | TokenKind::Neq
                | TokenKind::Lt
                | TokenKind::Le
                | TokenKind::Gt
                | TokenKind::Ge
                | TokenKind::And
                | TokenKind::Or
        )
    }

    /// 检查是否在函数参数上下文中
    fn is_in_func_param_context(&self) -> bool {
        // 简单实现：总是返回false，因为我们没有上下文跟踪
        // 在实际实现中，我们需要跟踪格式化器的状态
        false
    }

    /// 格式化函数参数（不添加换行）
    fn format_func_param(&mut self, children: &[AstNode]) -> Result<(), String> {
        // 格式化参数但不添加换行
        for (i, child) in children.iter().enumerate() {
            self.format_node(child)?;
            
            // 在int后加空格
            if i == 0 {
                self.add_space();
            }
        }
        Ok(())
    }
}

/// 后处理格式化输出，应用额外的格式化规则
fn postprocess_formatting(output: &str) -> String {
    let mut result = String::new();
    let mut chars = output.chars().peekable();
    let mut indent_level = 0;
    
    while let Some(c) = chars.next() {
        result.push(c);
        
        // 跟踪缩进级别
        if c == '{' {
            indent_level += 1;
        } else if c == '}' {
            indent_level -= 1;
        }
        
        // 处理分号后换行
        if c == ';' {
            if let Some(&next_char) = chars.peek() {
                if next_char != '\n' {
                    result.push('\n');
                    // 添加正确的缩进
                    for _ in 0..indent_level {
                        result.push_str("    ");
                    }
                }
            } else {
                result.push('\n');
            }
            continue;
        }
        
        // 处理函数名和括号之间的空格
        if c == ')' {
            if let Some(&next_char) = chars.peek() {
                if next_char == '{' {
                    result.push(' ');
                }
            }
        }
        
        // 处理else语句
        if result.ends_with("else") {
            // 确保else后面有空格
            if let Some(&next_char) = chars.peek() {
                if next_char != ' ' && next_char != '{' && next_char != '\n' {
                    result.push(' ');
                }
            }
        }
        
        // 处理else if语句
        if result.ends_with("else if") {
            // 确保else if后面有空格
            if let Some(&next_char) = chars.peek() {
                if next_char != ' ' && next_char != '(' {
                    result.push(' ');
                }
            }
        }
        
        // 处理运算符空格
        if c == '=' {
            // 检查是否是 == 运算符
            if let Some(&next_char) = chars.peek() {
                if next_char == '=' {
                    // 处理 == 运算符
                    result.push('=');
                    chars.next(); // 消耗第二个 =
                    
                    // 确保 == 两边有空格
                    if let Some(prev_char) = result.chars().rev().nth(1) {
                        if prev_char != ' ' && prev_char != '(' && prev_char != '!' {
                            result.insert(result.len() - 2, ' ');
                        }
                    }
                    
                    // 确保 == 后面有空格
                    if let Some(&next_next_char) = chars.peek() {
                        if next_next_char != ' ' && next_next_char != ')' && next_next_char != ';' {
                            result.push(' ');
                        }
                    }
                } else {
                    // 处理 = 赋值运算符
                    // 确保 = 前面有空格（除非前面已经是空格或括号）
                    if let Some(prev_char) = result.chars().rev().nth(1) {
                        if prev_char != ' ' && prev_char != '(' {
                            result.insert(result.len() - 1, ' ');
                        }
                    }
                    
                    // 确保 = 后面有空格
                    if let Some(&next_char) = chars.peek() {
                        if next_char != ' ' && next_char != ';' {
                            result.push(' ');
                        }
                    }
                }
            }
        }
        
        // 处理语句之间的换行
        if c == '}' {
            if let Some(&next_char) = chars.peek() {
                // 确保}后面有换行（除非后面是else或者已经有换行）
                if next_char != '\n'  {
                    result.push('\n');
                    // 添加正确的缩进
                    for _ in 0..indent_level {
                        result.push_str("    ");
                    }
                }
            } else {
                // 如果是文件末尾，添加换行
                result.push('\n');
            }
        }
    }
    
    // 最后处理函数之间的空行
    let lines: Vec<&str> = result.lines().collect();
    let mut final_result = String::new();
    let mut prev_was_func = false;
    
    for (i, line) in lines.iter().enumerate() {
        let trimmed = line.trim();
        
        // 检查是否是函数定义
        let is_func_def = trimmed.starts_with("int ") && trimmed.contains('(') && 
                         !trimmed.contains(';') && trimmed.ends_with('{');
        let is_void_func = trimmed.starts_with("void ") && trimmed.contains('(') && 
                          !trimmed.contains(';') && trimmed.ends_with('{');
        
        if (is_func_def || is_void_func) && i > 0 && prev_was_func {
            final_result.push('\n');
        }
        
        final_result.push_str(line);
        if i < lines.len() - 1 {
            final_result.push('\n');
        }
        
        prev_was_func = is_func_def || is_void_func;
    }
    
    final_result
}

/// 格式化AST入口函数
pub fn format_ast(ast: &AstNode) -> Result<String, String> {
    let mut formatter = Formatter::new();
    formatter.format_comp_unit(ast)?;
    let processed_output = postprocess_formatting(&formatter.output);
    Ok(processed_output)
}
